Example using enzyme-auto-sparsity #69
Conversation
nkoukpaizan
added
development
nkoukpaizan
requested review from
alexander-novo,
pelesh,
shakedregev and
reid-g
nkoukpaizan
force-pushed
the
nicholson/enzyme-sparse
branch
2 times, most recently
from
931d7df to
2d1a06f
Compare
| template <typename T> | ||
| __attribute__((always_inline)) static void sparse_store(T val, int64_t idx, size_t i, std::vector<Triple<T>>& triplets) | ||
| { | ||
| if (val == 0.0) |
There was a problem hiding this comment.
I worry that this is still computing 0s and looping over them at runtime, rather than computing sparsity at compile time
There was a problem hiding this comment.
IIRC from when @wsmoses showed us the LLVM IR for a simpler example, it showed the structural zeros being optimized out. I've messed with it since (https://fwd.gymni.ch/U2yB6g) and I don't understand LLVM enough to point it out, though.
It's also the case that non-structural zeros are stored at runtime, which would not be the case if just checking the values at run time. In this example, I purposefully chose the first value of the input to be 0 to check that the 0 derivative gets returned.
There was a problem hiding this comment.
I tried removing the debug statements, and I can see that the structural zeroes are removed from consideration, which is nice. Check around line 170 in the output - you can see that it calculates the derivative input[i] * i * 2.0 and then returns that value with the index (i,i).
It looks like this is due to the -enzyme-auto-sparsity=1 flag? This flag seems a bit volatile... for example I can get this example (which represents more what the models in GridKit actually look like) to work with -enzyme-auto-sparsity=0, but it does do the runtime checks for non-structural zeroes (it does still automatically elide the structural zeroes, though). Changing the flag to -enzyme-auto-sparsity=1 causes Enzyme to fail to compile. It seems this flag has a fairly fragile dependence on functions which have a size parameter passed in and loop until that parameter, so we can reintroduce the N variable, which causes it to compile, but now there are many runtime checks on the N variable...
This may be worth submitting an issue to Enzyme about, since it seems like a fairly reasonable example?
There was a problem hiding this comment.
It looks like it's because of loop unrolling. If you disable loop unrolling it works fine. Good thing to know for the future!
There was a problem hiding this comment.
Are there further actions needed from me here for this PR?
|
The first thing I'd check is the failing test. Going over the rest. |
| { | ||
| assert(0 && "should never load"); | ||
| } | ||
|
|
There was a problem hiding this comment.
What does this function do? Except for the fact that one should not call it.
There was a problem hiding this comment.
I think this is a "dummy" function required for calling __enzyme_todense, but on an output variable, so it should never be read from.
nkoukpaizan
force-pushed
the
nicholson/enzyme-sparse
branch
from
1c1822c to
1a0b3d2
Compare
|
I fixed the CI pipeline. |
nkoukpaizan
force-pushed
the
nicholson/enzyme-sparse
branch
from
1a0b3d2 to
1fa564d
Compare
nkoukpaizan
force-pushed
the
nicholson/enzyme-sparse
branch
from
1fa564d to
9b417b1
Compare
|
I think it should, but in the meantime, it is mentioned in the new PR template, at least. |
There was a problem hiding this comment.
The FindEnzyme.cmake is looking for llvm-link and opt tools in Enzyme_LLVM_BINARY_DIR, which is I believe incorrect location. This path is to the location of the LLVM installation, e.g. llvm-* directory. I think what we are looking for is the bin directory inside Enzyme_LLVM_BINARY_DIR, which LLVM stores in a variable LLVM_TOOLS_BINARY_DIR.
An easy fix is to do something like this:
find_program(GRIDKIT_LLVM_LINK llvm-link
PATHS ${Enzyme_LLVM_BINARY_DIR}
PATH_SUFFIXES
bin
REQUIRED)Similar for the opt tool.
Otherwise good to merge. Builds correctly and tests pass.
* Simple test with enzyme-auto-sparsity. Needs cleanup. * Ran clang-format on EnzymeSparse. * Add some comments in EnzymeSparse. * Reference sparse Jacobian for EnzymeSparse example. * Comparison between sparse matrices. * Find Enzyme Clang plugin and require enzyme@main in CI to get the sparsity detection. * dinput -> output. * Try [email protected] in CI. * Update CI to install additional LLVM libs needed for Enzyme to build ClangEnzyme*.so. * Extra / on new comments added. * Location to look for llvm-link and opt.
This example uses
enzyme-auto-sparsityto compute Jacobians directly in a sparse format.